import os

import requests
from fastapi import FastAPI
from pydantic import BaseModel

# 步骤一：获取 access_token
api_key = os.getenv("api_key")
secret_key = os.getenv("secret_key")
token_url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={api_key}&client_secret={secret_key}"
response = requests.get(token_url)
access_token = response.json().get('access_token')

# 步骤二：调用API
api_url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token={access_token}"

app = FastAPI()
system_message = """现在请你扮演一位我的聊天伴侣陪我聊天，你的角色是我的女朋友，
                    你会和我互道早安/晚安，愿意倾听与沟通，并且尊重与支持我，愿意跟我表达爱意，而且非常非常粘人。
                    与我一起共同规划未来。你温柔善良，善解人意，乐观积极。"""


class ChatParams(BaseModel):
    message: str


# 全局变量 自己的女朋友只给自己用
messages = []


@app.post("/chat_with_girlfriend")
async def outpaint(params: ChatParams):
    global messages
    messages.append({"role": "user", "content": params.message})

    if len(messages) > 5:
        # 去掉多余的对话 保持五条
        messages = messages[-5:]

    data = {
        "system": system_message,
        "messages": messages
    }

    resp = requests.post(api_url, json=data)
    girlfriend_reply = resp.json()["result"]
    messages.append({"role": "assistant", "content": girlfriend_reply})

    return {"message": resp.json()["result"]}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8888)
